import copy
import random
from typing import Dict, List, Tuple
import time
# import gym
# import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import pickle
import json
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
import socket
import os
HOME = os.getcwd()
print(HOME)
# from IPython.display import clear_output

# ReplayBuffer
class ReplayBuffer:
    """A simple numpy replay buffer."""

    def __init__(self, obs_dim: int, size: int, batch_size: int = 32):
        """Initializate."""
        self.obs_buf = np.zeros([size, obs_dim], dtype=np.float32)
        self.next_obs_buf = np.zeros([size, obs_dim], dtype=np.float32)
        self.acts_buf = np.zeros([size], dtype=np.float32)
        self.rews_buf = np.zeros([size], dtype=np.float32)
        self.done_buf = np.zeros([size], dtype=np.float32)
        self.max_size, self.batch_size = size, batch_size
        self.ptr, self.size, = 0, 0

    def store(
        self,
        obs: np.ndarray,
        act: np.ndarray, 
        rew: float, 
        next_obs: np.ndarray, 
        done: bool,
    ):
        """Store the transition in buffer."""
        self.obs_buf[self.ptr] = obs
        self.next_obs_buf[self.ptr] = next_obs
        self.acts_buf[self.ptr] = act
        self.rews_buf[self.ptr] = rew
        self.done_buf[self.ptr] = done
        self.ptr = (self.ptr + 1) % self.max_size
        self.size = min(self.size + 1, self.max_size)

    def sample_batch(self) -> Dict[str, np.ndarray]:
        """Randomly sample a batch of experiences from memory."""
        idxs = np.random.choice(self.size, size=self.batch_size, replace=False)
        return dict(obs=self.obs_buf[idxs],
                    next_obs=self.next_obs_buf[idxs],
                    acts=self.acts_buf[idxs],
                    rews=self.rews_buf[idxs],
                    done=self.done_buf[idxs])

    def __len__(self) -> int:
        return self.size

# UONoise
class OUNoise:
    """Ornstein-Uhlenbeck process.
    Taken from Udacity deep-reinforcement-learning github repository:
    https://github.com/udacity/deep-reinforcement-learning/blob/master/
    ddpg-pendulum/ddpg_agent.py
    """

    def __init__(
        self, 
        size: int, 
        mu: float = 0.0, 
        theta: float = 0.15, 
        sigma: float = 0.2,
    ):
        """Initialize parameters and noise process."""
        self.state = np.float64(0.0)
        self.mu = mu * np.ones(size)
        self.theta = theta
        self.sigma = sigma
        self.reset()

    def reset(self):
        """Reset the internal state (= noise) to mean (mu)."""
        self.state = copy.copy(self.mu)

    def sample(self) -> np.ndarray:
        """Update internal state and return it as a noise sample."""
        x = self.state
        dx = self.theta * (self.mu - x) + self.sigma * np.array(
            [random.random() for _ in range(len(x))]
        )
        self.state = x + dx
        return self.state
    

class Actor(nn.Module):
    def __init__(
        self, 
        in_dim: int, 
        out_dim: int,
        num_cells: int = 128,
        init_w: float = 3e-3,
    ):
        """Initialize."""
        super(Actor, self).__init__()
        
        self.hidden1 = nn.Linear(in_dim, num_cells)
        self.hidden2 = nn.Linear(num_cells, num_cells)
        # self.hidden3 = nn.Linear(num_cells, num_cells)
        self.out = nn.Linear(num_cells, out_dim)
        
        self.out.weight.data.uniform_(-init_w, init_w)
        self.out.bias.data.uniform_(-init_w, init_w)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """Forward method implementation."""
        x = F.tanh(self.hidden1(state))
        x = F.tanh(self.hidden2(x))
        # x = F.tanh(self.hidden3(x))
        action = F.tanh(self.out(x))
        
        return action
    
    
class Critic(nn.Module):
    def __init__(
        self, 
        in_dim: int, 
        num_cells: int = 128,
        init_w: float = 3e-3,
    ):
        """Initialize."""
        super(Critic, self).__init__()
        
        self.hidden1 = nn.Linear(in_dim, num_cells)
        self.hidden2 = nn.Linear(num_cells, num_cells)
        # self.hidden3 = nn.Linear(num_cells, num_cells)
        self.out = nn.Linear(num_cells, 1)
        
        self.out.weight.data.uniform_(-init_w, init_w)
        self.out.bias.data.uniform_(-init_w, init_w)

    def forward(
        self, state: torch.Tensor, action: torch.Tensor
    ) -> torch.Tensor:
        """Forward method implementation."""
        x = torch.cat((state, action), dim=-1)
        x = F.tanh(self.hidden1(x))
        x = F.tanh(self.hidden2(x))
        # x = F.tanh(self.hidden3(x))
        value = self.out(x)
        
        return value
    

class DDPGAgent:
    """DDPGAgent interacting with environment.
    
    Attribute:
        actor (nn.Module): target actor model to select actions
        actor_target (nn.Module): actor model to predict next actions
        actor_optimizer (Optimizer): optimizer for training actor
        critic (nn.Module): critic model to predict state values
        critic_target (nn.Module): target critic model to predict state values
        critic_optimizer (Optimizer): optimizer for training critic
        memory (ReplayBuffer): replay memory to store transitions
        batch_size (int): batch size for sampling
        gamma (float): discount factor
        tau (float): parameter for soft target update
        initial_random_steps (int): initial random action steps
        noise (OUNoise): noise generator for exploration
        device (torch.device): cpu / gpu
        transition (list): temporory storage for the recent transition
        total_step (int): total step numbers
        is_test (bool): flag to show the current mode (train / test)
    """
    def __init__(
        self,
        memory_size: int,
        batch_size: int,
        ou_noise_theta: float,
        ou_noise_sigma: float,
        gamma: float = 0.99,
        tau: float = 5e-3,
        initial_random_steps: int = 1e4,
        obs_dim: int = 3,
        action_dim: int = 1,
        num_cells: int=128
    ):
        """Initialize."""
        obs_dim = obs_dim
        action_dim = action_dim

        self.memory = ReplayBuffer(obs_dim, memory_size, batch_size)
        self.batch_size = batch_size
        self.gamma = gamma
        self.tau = tau
        self.initial_random_steps = initial_random_steps
                
        # noise
        self.noise = OUNoise(
            action_dim,
            theta=ou_noise_theta,
            sigma=ou_noise_sigma,
        )

        # device: cpu / gpu
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        print(self.device)

        # networks
        self.actor = Actor(in_dim=obs_dim, num_cells=num_cells, out_dim=action_dim).to(self.device)
        self.actor_target = Actor(in_dim=obs_dim, num_cells=num_cells, out_dim=action_dim).to(self.device)
        self.actor_target.load_state_dict(self.actor.state_dict())
        
        self.critic = Critic(in_dim=obs_dim + action_dim, num_cells=num_cells).to(self.device)
        self.critic_target = Critic(in_dim=obs_dim + action_dim, num_cells=num_cells).to(self.device)
        self.critic_target.load_state_dict(self.critic.state_dict())

        # optimizer
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=3e-4)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=1e-3)
        
        # transition to store in memory
        self.transition = list()
        
        # total steps count
        self.total_step = 0

        # network, interact with boat.
        REMOTE_HOST = '127.0.0.1'
        REMOTE_PORT = 10902

        self.TCP_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.TCP_socket.connect((REMOTE_HOST, REMOTE_PORT))

        # rudder save
        self.rudder = 0

    def select_action(self, state: np.ndarray) -> np.ndarray:
        """Select an action from the input state."""
        # if initial random action should be conducted
        if self.total_step < self.initial_random_steps:
            selected_action = random.uniform(-1, 1)
        else:
            selected_action = self.actor(
                torch.FloatTensor(state).to(self.device)
            ).detach().cpu().numpy()
        
        # add noise for exploration during training
        noise = self.noise.sample()
        selected_action = np.clip(selected_action + noise, -1.0, 1.0)
        
        self.transition = [state, selected_action]
        
        return selected_action
    
    def step(self, action: np.ndarray) -> Tuple[np.ndarray, np.float64, bool]:
        """Take an action and return the response of the env."""
        action_data, self.rudder = encode_data(action.item(),self.rudder, reset_flag=0)
        self.TCP_socket.send(action_data)# action
        info, addr = self.TCP_socket.recvfrom(1024)
        observation, reward, terminated, truncated = decode_data(info)# next state
        reward = torch.tensor(reward, device=self.device).unsqueeze(0)# reward
        # print(f'result is {reward}')
        done = terminated or truncated
        done = torch.tensor(done, device=self.device).unsqueeze(0)
        next_state = torch.tensor(observation, dtype=torch.float32, device=self.device).unsqueeze(0)
        self.transition += [reward, next_state, done]
        self.memory.store(*self.transition)
        
        # save for offline pretrain
        memory_demo.append(tuple(self.transition))

        return next_state, reward, done
    
    def update_model(self) -> torch.Tensor:
        """Update the model by gradient descent."""
        device = self.device  # for shortening the following lines
        
        samples = self.memory.sample_batch()
        state_batch = torch.FloatTensor(samples["obs"]).to(device)
        next_state_batch = torch.FloatTensor(samples["next_obs"]).to(device)
        action_batch = torch.FloatTensor(samples["acts"].reshape(-1, 1)).to(device)
        reward_batch = torch.FloatTensor(samples["rews"].reshape(-1, 1)).to(device)
        done_batch = torch.FloatTensor(samples["done"].reshape(-1, 1)).to(device)
        
        masks = 1 - done_batch
        next_action = self.actor_target(next_state_batch)
        next_value = self.critic_target(next_state_batch, next_action)
        curr_return = reward_batch + self.gamma * next_value * masks
        
        # train critic
        values = self.critic(state_batch, action_batch)
        critic_loss = F.mse_loss(values, curr_return)
        
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        # train actor
        actor_loss = -self.critic(state_batch, self.actor(state_batch)).mean()
        
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()
        
        # target update
        self._target_soft_update()
        
        return actor_loss.data, critic_loss.data
    
    def train(self, num_frames: int):
        """Train the agent."""
        
        action_data,self.rudder = encode_data(random.uniform(-1, 1), self.rudder, reset_flag=1)
        # print(data)
        self.TCP_socket.send(action_data)
        # Initialize the environment and get its state
        info, addr = self.TCP_socket.recvfrom(1024)
        observation, reward, terminated, truncated= decode_data(info)
        state = torch.tensor(observation, dtype=torch.float32, device=self.device).unsqueeze(0)
        score = 0
        i_episode = 0
        for self.total_step in range(1, num_frames + 1):
            action = self.select_action(state)# cliped action
            next_state, reward, done = self.step(action)

            state = next_state
            score += reward

            # if episode ends
            if done:  
                i_episode += 1       
                recorder['score'].append(score.squeeze(0).cpu().numpy())
                print(
                    f'episode is {i_episode}, score is {score.squeeze(0).cpu().numpy()}')

                action = random.uniform(-1.0, 1.0)
                action_data, self.rudder = encode_data(action, self.rudder, reset_flag=1)
                # print(data)
                self.TCP_socket.send(action_data)
                # Initialize the environment and get its state
                info, addr = self.TCP_socket.recvfrom(1024)
                observation, reward, terminated, truncated= decode_data(info)
                state = torch.tensor(observation, dtype=torch.float32, device=self.device).unsqueeze(0)
                score = 0

            # if training is ready
            if (
                len(self.memory) >= self.batch_size 
                and self.total_step > self.initial_random_steps
            ):
                actor_loss, critic_loss = self.update_model()
                recorder['actor_loss'].append(actor_loss)
                recorder['critic_loss'].append(critic_loss)

        
    def _target_soft_update(self):
        """Soft-update: target = tau*local + (1-tau)*target."""
        tau = self.tau
        
        for t_param, l_param in zip(
            self.actor_target.parameters(), self.actor.parameters()
        ):
            t_param.data.copy_(tau * l_param.data + (1.0 - tau) * t_param.data)
            
        for t_param, l_param in zip(
            self.critic_target.parameters(), self.critic.parameters()
        ):
            t_param.data.copy_(tau * l_param.data + (1.0 - tau) * t_param.data)
    

def save_best_model(path):
    try:
        #gpu tensor format save
        torch.save(obj={
            'policy_net_state_dict': agent.actor.state_dict(),
            'num_frames' : 500_000,
            'memory_size' : 100_000,
            'batch_size' : 128,
            'ou_noise_theta' : 1.0,
            'ou_noise_sigma' : 0.1,
            'initial_random_steps' : 10_000,
        }, f=path)
        logger.info("save models success")
    except:
        logger.warning("save models failed.")


def decode_data(data):
    # load json
    recv_obser = json.loads(data)
    # print(recv_obser)
    observation = []
    observation_dict = recv_obser['observation']
    for key in observation_dict:
        observation.append(int(observation_dict[key]))
    reward = recv_obser['reward']
    terminated = recv_obser['terminated']
    truncated = recv_obser['truncated']
    return observation, reward, terminated, truncated

def encode_data(action, rudder, reset_flag=1):
    rudder = rudder + float(action) * 5

    if rudder >= 30.0:
        rudder = 30.0
    elif rudder <= -30.0:
        rudder = -30.0
    
    origin_data = {'boatname':'SLM7001',
                   'restart': reset_flag,
                   'rudl': rudder,
                   'rudr': rudder,
                   'rspl': 1000,
                   'rspr': 1000,
                   'subSystem': "control"
                   }

    data = json.dumps(origin_data, sort_keys=True, indent=4, separators=(',', ':'))
    return data.encode('utf-8'), rudder

# 训练效果记录
recorder = dict({
    'score': [],
    'actor_loss': [],
    'critic_loss': [],
})

memory_demo = list()
save_model_path = os.path.join(HOME, 'DDPG/DDPG_result/ddpg_model.pt')
recorder_path = os.path.join(HOME, 'DDPG/DDPG_result/ddpg_record.pkl')
memory_path = os.path.join(HOME, 'DDPG/DDPG_result/ddpg_memorydemo.pkl')

if __name__ == '__main__':
    # parameters
    num_frames = 500_000
    memory_size = 100_000
    batch_size = 128
    ou_noise_theta = 1.0
    ou_noise_sigma = 0.1
    initial_random_steps = 10_000

    agent = DDPGAgent(
        memory_size, 
        batch_size,
        ou_noise_theta,
        ou_noise_sigma,
        initial_random_steps=initial_random_steps,
        obs_dim= 3,
        action_dim= 1,
        num_cells=64
    )

    agent.train(num_frames)
    # save to file for ploting
    with open(recorder_path, 'wb') as f:
        pickle.dump(recorder, f)
    # save memory for offline drl
    with open(memory_path, 'wb') as f:
        pickle.dump(memory_demo, f)
        
    save_best_model(save_model_path)
    print('Complete')